import os

import dnnlib
import dnnlib.tflib as tflib
import numpy as np
from PIL import Image

import tensorflow as tf


def initialize_tf():
    """
    初始化 TensorFlow
    """
    tflib.init_tf()

def load_generator(model_path):
    """
    加载预训练的 StyleGAN2 生成器
    """
    with dnnlib.util.open_url(model_path) as f:
        _G, _D, Gs = pickle.load(f)
    return Gs

def generate_images(Gs, num_images=1, truncation_psi=0.7, seed=None):
    """
    生成图像
    
    参数:
        Gs: 预训练的生成器网络
        num_images: 要生成的图像数量
        truncation_psi: truncation 参数
        seed: 随机种子
    """
    if seed is not None:
        np.random.seed(seed)

    # 生成随机潜码
    latent_size = Gs.input_shapes[0][1]
    latents = np.random.randn(num_images, latent_size)
    
    # 生成图像
    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    images = Gs.run(latents, None, truncation_psi=truncation_psi, randomize_noise=True, output_transform=fmt)
    
    return images

def save_images(images, output_dir='generated_images'):
    """
    保存生成的图像
    """
    os.makedirs(output_dir, exist_ok=True)
    for i, img in enumerate(images):
        Image.fromarray(img).save(f'{output_dir}/generated_{i:04d}.png')

def main():
    # 初始化 TensorFlow
    initialize_tf()
    
    # 模型路径（这里使用 FFHQ 预训练模型作为示例）
    model_path = 'gdrive:networks/stylegan2-ffhq-config-f.pkl'
    
    # 加载生成器
    print('加载预训练模型...')
    Gs = load_generator(model_path)
    
    # 生成参数设置
    num_images = 3
    truncation_psi = 0.7
    seed = 1000
    
    # 生成图像
    print('生成图像...')
    images = generate_images(Gs, 
                           num_images=num_images, 
                           truncation_psi=truncation_psi,
                           seed=seed)
    
    # 保存图像
    print('保存图像...')
    save_images(images)
    
    print('完成！')

if __name__ == '__main__':
    main()